"""Author: Sid Black (@sdtblck)"""
from googletrans import Translator
import json
from tqdm import tqdm
from multiprocessing import Pool, cpu_count
from functools import partial
import os
import random

translator = Translator(service_urls=[
        'translate.google.co.uk',
        'translate.google.de',
        'translate.google.fr',
        'translate.google.it',
        ])

countries = ['ukraine', 'united_kingdom']


def dump_jsonl(data, output_path, append=False):
    """
    Write list of objects to a JSON lines file.
    """
    mode = 'a+' if append else 'w'
    with open(output_path, mode, encoding='utf-8') as f:
        for line in data:
            json_record = json.dumps(line, ensure_ascii=False)
            f.write(json_record + '\n')
    print('Wrote {} records to {}'.format(len(data), output_path))


def _trans(i, dest, retries=0):
    if i:
        d = json.loads(i)
        try:
            translated = translator.translate(d['text'], src='en', dest=dest)
        except Exception as e:
            print(e)
            print(f'Error - retrying (attempt {retries})')
            if retries > 5:
                os.system(f'nordvpn c {random.choice(countries)}')
            return _trans(i, dest, retries=retries+1)
        return {'text': translated.text}


def translate_lambada(dest='fr'):
    results = []
    translate_fn = partial(_trans, dest=dest)
    with Pool(cpu_count()*4) as p:
        with open('lambada_test_en.jsonl', 'r') as f:
            f = list(f)
            for res in tqdm(p.imap_unordered(translate_fn, f)):
                results.append(res)

    return results

if __name__ == "__main__":
    de = translate_lambada(dest='de')
    dump_jsonl(de, 'lambada_test_de.jsonl')
